#include "clipper.hpp"
#include "nn_utils.h"
#include "paddleocr_post_process.h"
#ifdef LYNXI_PLUGIN
#include "lyn_plugin_dev.h"
#endif

namespace ocr {

float calculateAverage(const std::vector<float> &vec) {
  float sum = 0.0f;
  for (const auto &element : vec) {
    sum += element;
  }
  return sum / static_cast<float>(vec.size());
}

int lynPaddleOcrDetectPostProcess(PaddleOcrDetectPostProcessInfo_t *post_info) {
  int16_t *output_tensor = nullptr;
  lynBoxesInfo *boxesInfo = nullptr;
  int16_t *image_data = nullptr;
  int16_t *crop_images_data = nullptr;
  cropImagesInfo *cropsInfo = nullptr;
#ifdef LYNXI_PLUGIN
  output_tensor = (int16_t *)lynPluginGetVirtAddr(post_info->output_tensor);
  if (output_tensor == nullptr) {
    LOG_PLUGIN_E("get output tensor addr error\n");
    return -1;
  }

  boxesInfo = (lynBoxesInfo *)lynPluginGetVirtAddr(post_info->boxesInfo);
  if (boxesInfo == nullptr) {
    LOG_PLUGIN_E("get boxesInfo addr error\n");
    return -1;
  }

  image_data = (int16_t *)lynPluginGetVirtAddr(post_info->image_data);
  if (image_data == nullptr) {
    LOG_PLUGIN_E("get image_data addr error\n");
    return -1;
  }

  crop_images_data =
      (int16_t *)lynPluginGetVirtAddr(post_info->crop_images_data);
  if (crop_images_data == nullptr) {
    LOG_PLUGIN_E("get crop_images_data addr error\n");
    return -1;
  }

  cropsInfo = (cropImagesInfo *)lynPluginGetVirtAddr(post_info->cropsInfo);
  if (cropsInfo == nullptr) {
    LOG_PLUGIN_E("get cropsInfo addr error\n");
    return -1;
  }
#else
  output_tensor = (int16_t *)post_info->output_tensor;
  boxesInfo = post_info->boxesInfo;
  image_data = (int16_t *)post_info->image_data;
  crop_images_data = (int16_t *)post_info->crop_images_data;
  cropsInfo = post_info->cropsInfo;
#endif
  std::vector<TextBox> postprocess_rst = det_postprocess(
      output_tensor, post_info->img_width, post_info->img_height,
      post_info->m_width, post_info->m_height);

  boxesInfo->boxesNum = postprocess_rst.size();
  for (int i = 0; i < boxesInfo->boxesNum; i++) {
    std::vector<cv::Point> box = postprocess_rst[i].boxPoint;

    int collectX[4] = {box[0].x, box[1].x, box[2].x, box[3].x};
    int collectY[4] = {box[0].y, box[1].y, box[2].y, box[3].y};

    boxesInfo->boxes[i].xmax = int(*std::max_element(collectX, collectX + 4));
    boxesInfo->boxes[i].xmin = int(*std::min_element(collectX, collectX + 4));
    boxesInfo->boxes[i].ymax = int(*std::max_element(collectY, collectY + 4));
    boxesInfo->boxes[i].ymin = int(*std::min_element(collectY, collectY + 4));
    boxesInfo->boxes[i].score = postprocess_rst[i].score;
  }
  std::vector<cv::Mat> crop_images = crop(
      image_data, post_info->img_width, post_info->img_height, postprocess_rst);
  size_t offset = 0;
  size_t index = 0;
  for (const auto &mat : crop_images) {
    size_t dataSize = mat.total() * mat.elemSize();
    std::memcpy(crop_images_data + offset, mat.data, dataSize);
    offset += dataSize / sizeof(int16_t);

    cropsInfo->images[index].w = mat.size().width;
    cropsInfo->images[index].h = mat.size().height;
    index += 1;
  }

  return 0;
}

int lynPaddleOcrRecogPostProcess(PaddleOcrRecogPostProcessInfo_t *post_info) {
  int16_t *output_tensor = nullptr;
  mKeys *mkeys = nullptr;
  lynBoxesInfo *boxesInfo = nullptr;
#ifdef LYNXI_PLUGIN
  output_tensor = (int16_t *)lynPluginGetVirtAddr(post_info->output_tensor);
  if (output_tensor == nullptr) {
    LOG_PLUGIN_E("get output tensor addr error\n");
    return -1;
  }

  boxesInfo = (lynBoxesInfo *)lynPluginGetVirtAddr(post_info->boxesInfo);
  if (boxesInfo == nullptr) {
    LOG_PLUGIN_E("get boxesInfo addr error\n");
    return -1;
  }

  mkeys = (mKeys *)lynPluginGetVirtAddr(post_info->mkeys);
  if (mkeys == nullptr) {
    LOG_PLUGIN_E("get mkeys addr error\n");
    return -1;
  }

#else
  output_tensor = (int16_t *)post_info->output_tensor;
  mkeys = post_info->mkeys;
  boxesInfo = post_info->boxesInfo;
#endif
  TextBox res = TextBox();
  scoreToTextLine(output_tensor, res, mkeys, post_info->m_height,
                  post_info->m_width);
  boxesInfo->boxes[post_info->boxIndex].score =
      calculateAverage(res.textLine.charScores);

  strncpy(boxesInfo->boxes[post_info->boxIndex].label,
          res.textLine.text.c_str(),
          sizeof(boxesInfo->boxes[post_info->boxIndex].label) - 1);
  return 0;
}

} // namespace ocr